import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import os
import seaborn as sns
sns.set_context("paper",rc={"lines.linewidth": 0.5})
# sns.set_style("whitegrid")
path = os.listdir(os.getcwd())
DATASETS = ['MUTAG']
MODELS = ['GAT','GCN']
for model in MODELS:
    plt.figure()
    for dataset in DATASETS:
        shuffledGatTrainAcc=[]
        shuffledGatTestAcc=[]
        for fileName in path:
            if fileName.count(model)!=0 and fileName.count('shuffle')!=0 and fileName.count(dataset)!=0:
                data = pd.read_csv(fileName,sep=' ',header = None )
                shuffledGatTrainAcc.append(data[0].values.tolist())
                shuffledGatTestAcc.append(data[1].values.tolist())
        if shuffledGatTestAcc:
            shuffledGatTrainAcc = np.array(shuffledGatTrainAcc)
            shuffledGatTestAcc = np.array(shuffledGatTestAcc)
            plt.plot(range(1, shuffledGatTrainAcc.shape[1] + 1), shuffledGatTrainAcc.mean(axis=0), shuffledGatTestAcc.mean(axis=0))
        naiveGatTrainAcc = []
        naiveGatTestAcc = []
        for fileName in path:
            if fileName.count('naive')!=0 and fileName.count(dataset)!=0 and fileName.count(model)!=0:
                data = pd.read_csv(fileName, sep=' ', header=None)
                plt.plot(range(1,len(data)+1),data[1])
        gatTrainAcc=[]
        gatTestAcc = []
        for fileName in path:
            if fileName.count(model)!=0 and fileName.count(dataset)!=0 and fileName.count('.txt')!=0 and fileName.count('shuffle')==0 and fileName.count('naive')==0:
                data = pd.read_csv(fileName, sep=' ', header=None)
                gatTrainAcc.append(data[0].values.tolist())
                gatTestAcc.append(data[1].values.tolist())
        if len(gatTestAcc):
            gatTrainAcc = np.array(gatTrainAcc)
            gatTestAcc = np.array(gatTestAcc)
        plt.plot(range(1,len(data)+1),gatTestAcc.mean(axis=0))
        plt.legend(['shuffledGDVtrainAccMean', 'shuffledGDVtestAccMean','naiveTestAcc','GDVtestAccMean'])
        plt.title(dataset+'+ ShuffledGDV+ '+model)
        sns.despine()
plt.show()
